import os

import PIL
import imageio
import numpy as np
import skimage
from matplotlib import pyplot as plt
import cv2 as cv
from PIL import Image
import os.path
import glob
import random
import math
import data_util


ARTEFACT = 'artifact'
NOISE = 'noise'
DETAIL_LESS = 'detail_less'

H = 0
P = 1
#添加伪影
def add_artifact(img, h_p = H, shuffle = 0, downsample_interval = 5, protected_center_radious = 0.1):

    #对参数预处理
    # if(shuffle == 0 and downsample_interval%2 == 0):
    #     downsample_interval = downsample_interval+1

    # 傅里叶变换，得到复数的数组
    f = np.fft.fft2(img)
    # 低频在左上角，放到中心位置，依然是复数
    fshift = np.fft.fftshift(f)

    #加椒盐噪声
    # fshift = addnoise(fshift)

    #随机隔行采样
    x,y = fshift.shape

    if h_p == H:
        center_start = (0.5 - protected_center_radious / 2) * x
        center_end = (0.5 + protected_center_radious / 2) * x
        for i in range(x):
            if i >center_start and i < center_end:
                continue
            #shuffle:是否随机
            if (shuffle== 0 and i%downsample_interval == 0) or (shuffle ==1 and random.randint(0,downsample_interval) == 0) :
                    fshift[i , 0:y]=0
    else:
        center_start = (0.5 - protected_center_radious / 2) * y
        center_end = (0.5 + protected_center_radious / 2) * y
        for i in range(y):
            if i >center_start and i < center_end:
                continue
            if (shuffle== 0 and i%downsample_interval == 0) or (shuffle ==1 and random.randint(0,downsample_interval) == 0) :
                fshift[0:x , i]=0

    # 傅里叶逆变换，将中心的低频图像再移动到左上角
    ishift = np.fft.ifftshift(fshift)
    # 然后进行逆变换,得到的是复数数组
    iomage = np.fft.ifft2(ishift)
    # 将复数数组去绝对值，转化为实数
    image = np.abs(iomage)
    return image

def circle_downsample(cur_img,dowsample_radius_ratio):
    fft = np.fft.fft2(cur_img)
    fshift = np.fft.fftshift(fft)
    downsample = mask_circle(fshift, dowsample_radius_ratio)
    ifshift = np.fft.ifftshift(downsample)
    ifft = np.fft.ifft2(ifshift)
    img = np.abs(ifft)
    return img


#保留圆心
def mask_circle(image , dowsample_radius_ratio):
    shape = image.shape
    d = dowsample_radius_ratio * max(shape)
    center_point = tuple(map(lambda x: (x-1)/2, shape))
    for i in range(shape[0]):
        for j in range(shape[1]):
            dis = cal_distance(center_point, (i, j))
            if dis > d:
                image[i, j] = 0
    return image


def cal_distance(pa, pb):
    # d值决定中心圆大小，遍历并计算频率域中的每一个点与这个圆边界的距离
    dis = math.sqrt((pa[0]-pb[0])**2+(pa[1]-pb[1])**2)
    return dis


#添加椒盐噪声
def addnoise(src):
    h,w = src.shape[:2]  #获取图像的宽高信息
    nums = 5000
    rows = np.random.randint(0, h, (5000), dtype = np.int)
    cols = np.random.randint(0,w,(5000),dtype = np.int)
    for i in range(nums):
        if i%2 == 1:
            src[rows[i],cols[i]] = 255#奇数则产生白色
        else:src[rows[i],cols[i]] =0 #偶数则产生黑椒盐色
    return src

#添加高斯噪声
def add_gause(src):
    copy = np.copy(src)
    gnoise = np.zeros(src.shape,src.dtype)
    m=16  #噪声均值
    s=31 #噪声方差
    cv.randn(gnoise,m,s)##产生高斯噪声
    dst = cv.add(copy,gnoise)#将高斯噪声图像加到原图上去
    return dst



def show_diffrent_downsample(path, mode=DETAIL_LESS, ratios=(0.05,0.06, 0.07, 0.08, 0.09, 0.1,0.12, 0.15),
                             h_p=H, shuffle=1, downsample_intervals=(3,4,5,2,3,3,4,5), protected_center_radious=(0.001,0.002,0.003,0.004,0.005,0.006,0.007,0.008),
                             noise_mode = "gaussian", gauss_max = 0.4, salt_amount_max = 0.3):

    cur_img = PIL.Image.open(path)
    cur_img = np.array(cur_img)
    length = len(ratios)
    h = math.sqrt(length + 1)
    if h % 1 == 0:
        h = int(h)
        p = int((length + 1) / h)
    else:
        h = int(h)
        p = int((length + 1) / h) + 1

    plt.subplot(int(str(h) + str(p) + str(1)))
    plt.imshow(cur_img, cmap="gray")
    plt.title(mode)
    plt.axis('off')

    for i in range(length):
        if mode == data_util.DETAIL_LESS:
            a = circle_downsample(cur_img, ratios[i])
            plt.subplot(int(str(h) + str(p) + str(i + 2)))
            plt.imshow(a, cmap="gray")
            plt.title('d'+str(ratios[i])[0:5])
            plt.axis('off')
        elif mode == data_util.ARTEFACT:
            a = add_artifact(cur_img, h_p=h_p, shuffle=shuffle, downsample_interval=downsample_intervals[i],
                                        protected_center_radious=protected_center_radious[i])
            plt.subplot(int(str(h) + str(p) + str(i + 2)))
            plt.imshow(a, cmap="gray")
            title = 'h:{}s:{},i:{},c:{}'.format(str(h_p),shuffle , str(downsample_intervals[i])[0:5], str(protected_center_radious[i])[0:5])
            plt.title(title)
            plt.axis('off')
        else:
            if(noise_mode == 'gaussian') :
                mean = random.uniform(0.01 , gauss_max)
                var = random.uniform(0.01 , gauss_max - mean)
                a = skimage.util.random_noise(cur_img, noise_mode, var=var,mean=mean)
                title = 'v:{},m:{}'.format(str(var)[0:5] , str(mean)[0:5])
            else:
                salt_amount =  random.uniform(0.01 , salt_amount_max)
                a = skimage.util.random_noise(cur_img, noise_mode, amount=salt_amount)
                title = 'a:{}'.format(str(salt_amount)[0:5])
            plt.subplot(int(str(h) + str(p) + str(i + 2)))
            plt.imshow(a, cmap="gray")

            plt.title(title)
            plt.axis('off')
    plt.show()


def rename(dir, out):
    # 原始图片路径
    dir = "E:\download\dataset\IXI-T1_output_1"
    out = "E:\download\dataset\IXI-T1_output"
    # 获取该路径下所有图片
    filelist = os.listdir(dir)
    a = 1
    for file in filelist:
        # 原始路径
        Olddir = os.path.join(dir, file)
        # filename = os.path.splitext(file)[0]
        # .bmp
        filetype = os.path.splitext(file)[1]
        # 需要存储的路径 a 是需要定义修改的文件名
        Newdir = os.path.join(out, str(a) + filetype)
        os.rename(Olddir, Newdir)
        a += 1


#添加线性掩膜
def get_linear_sampling_mask(imageshape, center_sample_ratio=0.3, edge_sample_ratio=0.5):
    '''
    生成线性降采样的掩膜
    :param imageshape: 图像大小，tuple或list，形如：(64,64)，表示图像的长宽分别为64
    :param center_sample_ratio: 中心采样的比例，默认为0.16，即采样比例为16%
    :param edge_sample_ratio: 边缘采样比例，默认为0.3，即采样比例为30%
    :return: 2D array，线性降采样的掩膜。被掩膜掉的部分为0，未被掩膜掉的部分为1
    '''
    if (type(imageshape)!=list) and (type(imageshape)!=tuple):
        raise Exception("imageshape参数必须是列表或者元组")

    if center_sample_ratio==1.0 or edge_sample_ratio==1.0:
        return np.ones(imageshape)

    mask_zero = np.zeros(imageshape)
    rows, cols = mask_zero.shape

    center_point = int(rows // 2 - 1)
    nb_center_lines = int(rows * center_sample_ratio)
    print(f'{nb_center_lines=}')
    center_from = center_point - nb_center_lines // 2
    center_to = center_point + nb_center_lines // 2

    mask_zero[center_from:center_to, :] = 1

    nb_high_edge = rows - center_to
    nb_low_edge = center_from

    nb_high_edge_keep = int(nb_high_edge * edge_sample_ratio)
    nb_low_edge_keep = int(nb_low_edge * edge_sample_ratio)
    print(f'{nb_high_edge_keep=}')
    print(f'{nb_low_edge_keep=}')
    nb_high_edge_space = int(nb_high_edge // nb_high_edge_keep)
    nb_low_edge_space = int(nb_low_edge // nb_low_edge_keep)
    space = min(nb_low_edge_space, nb_high_edge_space)

    print(f'{nb_low_edge_space=}', f'{nb_high_edge_space=}', f'{space=}')
    mask_zero[center_to:rows:space, :] = 1
    mask_zero[0:center_from:space, :] = 1

    return mask_zero



# 提取目录下所有图片,更改尺寸后保存到另一目录
#将jpgfile转换成目标像素并保存到outdir下,文件名为filename
def convertpng(jpgfile, outdir,filename, width=128, height=128):
    img = Image.open(jpgfile)
    try:
        new_img = img.resize((width, height), Image.BILINEAR)
        new_img.save(os.path.join(outdir, filename))
    except Exception as e:
        print(e)


def show_serveral_downsample_imgs(img):
    downsample_intervals = [random.randint(2,5) for i in range(8)]
    protected_center_radious = [random.uniform(0.01,0.1) for i in range(8)]
    show_diffrent_downsample(img, mode=ARTEFACT,
                             h_p=random.randint(0, 1), shuffle=random.randint(0, 1),
                             downsample_intervals=downsample_intervals,
                             protected_center_radious=protected_center_radious)
    ratios = [random.uniform(0.08,0.2) for i in range(8)]
    show_diffrent_downsample(img, mode=DETAIL_LESS , ratios=ratios)
    show_diffrent_downsample(img, mode=NOISE , gauss_max=0.35)
    show_diffrent_downsample(img, mode=NOISE, noise_mode='s&p',
                             salt_amount_max=0.35)

if __name__ == '__main__':
    img = r'E:/download/Dataset/keras/IXI-T2/output/1.png'
    show_serveral_downsample_imgs(img)
    img = r'E:/download/Dataset/keras/IXI-DTI/output/1.png'
    show_serveral_downsample_imgs(img)
    img = r'E:/download/Dataset/keras/IXI-T1/output/1.png'
    show_serveral_downsample_imgs(img)
    img = r'E:/download/Dataset/keras/IXI-pd/output/1.png'
    show_serveral_downsample_imgs(img)
